'''
This file contains classes for drawing EBL patterns in a variety of output
formats, including a matlab script which when run will generate a pdf
output. Note that all these are implicitly 2d, and units are nm.

The grid class generates a set of cross-shape markers, as used for the flake
location fields.

Created on 10 Dec 2010

@author: jcb91
'''
import copy
import math
import xml.dom.minidom  # For xml documents

class point:
    def __init__(self,*args):
        # for no input args, initialise to origin
        if len(args)==0:
            self.x=0
            self.y=0
            return
        # for initialising from a point (copy constructor)
        elif isinstance(args[0],point):
            self.x=args[0].x
            self.y=args[0].y
            return
        # for initialising from a list of numbers
        elif isinstance(args[0],list):
            self.x=args[0][0]
            self.y=args[0][1]
            return
        # for initialising from numbers (default)
        elif (isinstance(args[0],int)) or (isinstance(args[0],float)):
            self.x = args[0]
            self.y = args[1]
            return
        # if we got to here, we don't know what was passed in!
        else:
            raise TypeError('initialising a point with argument type '+str(type(args[0]))+' is not supported')
    def __str__(self):
        return "("+self.__repr__()+")"
    def __repr__(self):
        return str(self.x)+","+str(self.y)
    def __neg__(self):
        return point(-self.x,-self.y)
    def __add__(self, other):
        if not isinstance(other,point):
            return NotImplemented
        else:
            return point(self.x+other.x,self.y+other.y)
    def __sub__(self, other):
        if not isinstance(other,point):
            return NotImplemented
        else:
            return self+(-other)
    def __mul__(self, other):
        if not (isinstance(other,int) or isinstance(other,float)):
            return NotImplemented
        else:
            return point(self.x*other,self.y*other)
    def __rmul__(self, other):
        return self.__mul__(other)
    def __div__(self, other):
        return self.__mul__(1.0/other)
    def __iadd__(self, other):
        if not isinstance(other,point):
            return NotImplemented
        else:
            self.translate(other)
            return self
    def __isub__(self, other):
        if not isinstance(other,point):
            return NotImplemented
        else:
            self.translate(-other)
            return self
    def __imul__(self, other):
        if not (isinstance(other,int) or isinstance(other,float)):
            return NotImplemented
        else:
            self.x=self.x*other
            self.y=self.y*other
            return self
    def __idiv__(self, other):
        if not (isinstance(other,int) or isinstance(other,float)):
            return NotImplemented
        else:
            self.x=self.x/float(other)
            self.y=self.y/float(other)
            return self
    def __abs__(self):
        return pow(self.x**2+self.y**2,0.5)
    def __getitem__(self,key):
        if key==0:
            return self.x
        elif key==1:
            return self.y
        else:
            raise IndexError('point objects only have 2 elements')
    def translate(self,*args):
        if len(args)==0:
            return self
        # translate by a point
        elif isinstance(args[0],point):
            self.x += args[0].x
            self.y += args[0].y
            return self
        # translate by a list of numbers
        elif isinstance(args[0],list):
            self.x+=args[0][0]
            self.y+=args[0][1]
            return self
        # translate by 2 number arguments
        elif (isinstance(args[0],int)) or (isinstance(args[0],float)):
            self.x += args[0]
            self.y += args[1]
            return self
        # if we got to here, we don't know what was passed in!
        else:
            raise TypeError('translating a point with argument type '+str(type(args[0]))+' is not supported')
    def flipx(self):
        self.x = -self.x
        return self
    def flipy(self):
        self.y = -self.y
        return self

class shape:
    def __init__(self,*args):
        if isinstance(args[0],list):
            if isinstance(args[0][0],point) or isinstance(args[0][0],list):
                self.coords=[]
                for i in range(len(args[0])):
                    self.coords.append(point(args[0][i]))
        elif isinstance(args[0],shape):
            self.coords=copy.deepcopy(args[0].coords)
        self.vertices=len(self.coords)
    def __str__(self):
        return "<shape with "+str(self.vertices)+" vertices>"
    def __repr__(self):
        return self.__str__()
    def translate(self,*args):
        for i in range(len(self.coords)):
            self.coords[i].translate(*args)
        return self
    def flipx(self):
        for i in range(len(self.coords)):
            self.coords[i].flipx()
        return self
    def flipy(self):
        for i in range(len(self.coords)):
            self.coords[i].flipy()
        return self
    def x(self):
        retstr="";
        for i in range(len(self.coords)):
            retstr=retstr+str(self.coords[i].x)+", "
        retstr=retstr+str(self.coords[0].x)
        return retstr
    def y(self):
        retstr="";
        for i in range(len(self.coords)):
            retstr=retstr+str(self.coords[i].y)+", "
        retstr=retstr+str(self.coords[0].y)
        return retstr

class square(shape):
    def __init__(self,sidelength):
        s=sidelength/2
        coords=[[-s, s],
                [ s, s],
                [ s,-s],
                [-s,-s]]
        shape.__init__(self,coords)

class bar(shape):
    def __init__(self,width,height):
        self.w=width/2.0
        self.h=height/2.0
        coords=[[-self.w, self.h],
                [ self.w, self.h],
                [ self.w,-self.h],
                [-self.w,-self.h]]
        shape.__init__(self,coords)

class cross(shape):
    def __init__(self,linelength,linewidth,*drawarms):
        if len(drawarms)==0:
            drawarms=[True, True, True, True]
        else:
            drawarms=drawarms[0]
            if len(drawarms)<4:
                for i in range(len(drawarms),4):
                    drawarms.append(False)

        if not isinstance(linelength,list):
            linelength = [linelength, linelength, linelength, linelength]
        elif len(linelength)<4:
            tmplinelength = linelength
            linelength=[];
            j=0
            for i in range(4):
                if drawarms[i]:
                    linelength.append(tmplinelength[j])
                    j+=1
                else:
                    linelength.append(linewidth)

        if not isinstance(linewidth,list):
            linewidth = [linewidth, linewidth, linewidth, linewidth]
        elif len(linewidth)<4:
            tmplinewidth = linewidth
            j=0
            for i in range(4):
                if drawarms[i]:
                    linewidth.append(tmplinewidth[j])
                    j+=1
                else:
                    linewidth.append(linewidth)

        if not min(linelength)>max(linewidth)/2:
            raise Exception('linewidth must be < line length!')

        coords=[]
        if drawarms[0]:
            coords.append(point( linelength[0] ,-linewidth[0]/2))
            coords.append(point( linelength[0] , linewidth[0]/2))
        if drawarms[0] == drawarms[1]:
            coords.append(point( linewidth[1]/2, linewidth[0]/2))
        if drawarms[1]:
            coords.append(point( linewidth[1]/2, linelength[1] ))
            coords.append(point(-linewidth[1]/2, linelength[1] ))
        if drawarms[1] == drawarms[2]:
            coords.append(point(-linewidth[1]/2, linewidth[2]/2))
        if drawarms[2]:
            coords.append(point(-linelength[2] , linewidth[2]/2))
            coords.append(point(-linelength[2] ,-linewidth[2]/2))
        if drawarms[2] == drawarms[3]:
            coords.append(point(-linewidth[3]/2,-linewidth[2]/2))
        if drawarms[3]:
            coords.append(point(-linewidth[3]/2,-linelength[3] ))
            coords.append(point( linewidth[3]/2,-linelength[3] ))
        if drawarms[3] == drawarms[0]:
            coords.append(point(linewidth[3]/2,-linewidth[0]/2))
        shape.__init__(self,coords)

class field:
    def __init__(self,fieldsize):
        self.fieldsize  = fieldsize*1000
        self.dots       = 60000
        self.shapelists = []
    def writematlab(self,filename="gridpattern0"):
        f=open(filename+".m","w")
        f.write("clf;\n")
        f.write("scale=1000;\n");
        f.write("dx="+str(self.fieldsize)+"/scale;\n")
        f.write("dy=dx;\n")
        f.write("pad=5000/scale;\n")
        f.write("axis equal;\n")
        f.write("xlim([0 dx]+[-pad pad])\n")
        f.write("ylim([-dy 0]+[-pad pad])\n")
        f.write("line([0 dx dx 0 0],[0 0 -dy -dy 0],'color','red')\n")
        f.write("shapelists={...\n")
        for shapes in self.shapelists:
            f.write(" "*7+"{");
            for i in range(len(shapes)):
                if not i==0:
                    f.write(" "*8)
                f.write("["+shapes[i].x()+"; "+shapes[i].y()+"]/scale")
                if not i==len(shapes)-1:
                    f.write(", ...\n")
            f.write(" "*7+"};\n")
        f.write("};\n")
        f.write("for ii=1:length(shapelists)\n")
        f.write("for jj=1:length(shapes)\n")
        f.write("    line(shapelists{ii}{jj}(1,:),shapelists{ii}{jj}(2,:),'color','black');\n")
        f.write("end\n")
        f.write("xlabel('x / \\mum');\n")
        f.write("ylabel('y / \\mum');\n")
        f.write("print('-dpdf','"+filename+"');\n")
        f.close()
    def writemathematica(self,filename="gridpattern0"):
        f=open(filename+'.nb', 'w')
        f.write('SetDirectory[NotebookDirectory[]]')
        f.write('Export["'+filename+'.pdf",Show[')
        for shapes in self.shapelists:
            for sh in shapes:
                f.write('Graphics[Polygon[{');
                for i in range(len(sh.coords)):
                    f.write('{'+str(sh.coords[i].x)+','+str(-sh.coords[i].y)+'}')
                    if not i == len(sh.coords)-1:
                        f.write(',')
                f.write('}]],')
        f.write('ImageSize->2400]]')
        f.close()
    def writeccc(self,filename="gridpattern0"):
        #raise Exception("this method should not be trusted, as I'm not yet sure\n"+\
        #                "what the coordinate specifications DWSL and DWTZL denote.")
        s=self.fieldsize/self.dots # a scaling factor
        #write .CCC file
        f=open(filename+'.CCC', 'w')
        f.write("/*--- 1.CCC ---*/\n")
        f.write("/* CZ"+str(self.fieldsize/1000000.0)+","+str(self.dots)+" */\n")
        f.write("PATTERN\n")
        for shapes in self.shapelists:
            for sh in shapes:
                f.write("DWTZL(")
                for c in sh.coords:
                    f.write(str(c.x/s)+','+str(c.y/s)+',')
                f.write("1,0);3\n")
        f.write("!END")
        f.close()
        # write accompanying .CON file
        f=open(filename+'.CON','w')
        f.write("/*--- "+filename+".CON ---*/\n")
        f.write("CZ"+str(self.fieldsize/1000000.0)+","+str(self.dots)+";\n")
        f.write("PC"+filename+";\n")
        f.write("0.00000,0.00000;\n")
        f.write("PP"+filename+";\n")
        f.write("0.00000,0.00000;\n")
        f.write("!END")
    def writexml(self,filename="gridpattern0"):
        f=open(filename+'.xml', 'w')
        #opening parts of file
        f.write('<?xml version="1.0" encoding="Shift_JIS" ?>\n')
        f.write('<header versionnumber="3">\n')
        f.write('<biginlib lastmod="109-8-4 16:7:42" lastacc="109-8-4 16:7:42">\n')
        f.write('<libname>'+filename+' LEDIT 10.20.0</libname>\n')
        f.write('<units>0.000000,0.000000</units>\n')
        for ii in range(len(self.shapelists)):
            f.write('<bgnstr creation="108-8-27 16:48:35" lastmod="109-8-2 20:1:42">\n')
            f.write('<strname>set'+str(ii)+'</strname>\n')
            for sh in self.shapelists[ii]:
                f.write('<boundary>\n');
                f.write('<layer number="46">\n');
                xml_coords='<xy>'
                #get all the coordinates and format them properly.
                for c in sh.coords:
                    xml_coords=xml_coords+str.format("{0:d},{1:d}",int(c.x),int(c.y))+' '
                #must end with the beginning coordinate
                xml_coords=xml_coords+str.format("{0:d},{1:d}",int(sh.coords[0].x),int(sh.coords[0].y))+'</xy>\n'
                f.write(xml_coords);
                f.write('</layer>\n');
                f.write('</boundary>\n');
            f.write('</bgnstr>\n');
        #closing part of file
        f.write('<bgnstr creation="108-8-27 16:48:35" lastmod="109-8-4 16:7:42">\n')
        f.write('<strname>pattern</strname>\n')
        for ii in range(len(self.shapelists)):
            f.write('<sref>\n')
            f.write('<sname>set'+str(ii)+'</sname>\n')
            f.write('<strans reflection="0" absolutemagnification="0" absoluteangle="0" />\n')
            f.write('<xy>0,0</xy>\n')
            f.write('</sref>\n')
        f.write('</bgnstr>\n')
        f.write('</biginlib>\n')
        f.write('</header>\n')
        f.close()
    def writexmldom(self,filename="gridpatterndom"):
        doc = xml.dom.minidom.Document()
        header = doc.createElement('header')
        header.setAttribute('versionnumber','3')
        biginlib = doc.createElement('biginlib')
        biginlib.setAttribute('lastmod','109-8-4 16:7:42')
        biginlib.setAttribute('lastacc','109-8-4 16:7:42')

        doc.appendChild(header)
        header.appendChild(biginlib)

        libname = doc.createElement('libname')
        libname.appendChild(doc.createTextNode('scale1 LEDIT 10.20.0'))
        units = doc.createElement('units')
        units.appendChild(doc.createTextNode('0.000000,0.000000'))

        biginlib.appendChild(libname)
        biginlib.appendChild(units)

        for ii in range(len(self.shapelists)):
            bgnstr = doc.createElement('bgnstr')
            bgnstr.setAttribute('creation','108-8-27 16:48:35')
            bgnstr.setAttribute('lastmod','109-8-2 20:1:42')
            strname = doc.createElement('strname')
            strname.appendChild(doc.createTextNode('set'+str(ii)))
            bgnstr.appendChild(strname)

            for sh in self.shapelists[ii]:
                boundary = doc.createElement('boundary')
                layer = doc.createElement('layer')
                layer.setAttribute('number','46')
                xy = doc.createElement('xy')
                xml_coords=''
                for c in sh.coords:
                    xml_coords=xml_coords+str.format("{0:d},{1:d}",int(c.x),int(c.y))+' '
                #must end with the beginning coordinate
                xml_coords=xml_coords+str.format("{0:d},{1:d}",int(sh.coords[0].x),int(sh.coords[0].y))
                bgnstr.appendChild(boundary)
                boundary.appendChild(layer)
                layer.appendChild(xy)
                xy.appendChild(doc.createTextNode(xml_coords))

            biginlib.appendChild(bgnstr)

        #closing part of file
        bgnstr = doc.createElement('bgnstr')
        bgnstr.setAttribute('creation','108-8-27 16:48:35')
        bgnstr.setAttribute('lastmod','109-8-4 16:7:42')
        strname = doc.createElement('strname')
        strname.appendChild(doc.createTextNode('pattern'))
        bgnstr.appendChild(strname)
        for ii in range(len(self.shapelists)):
            sref = doc.createElement('sref')
            sname = doc.createElement('sname')
            sname.appendChild(doc.createTextNode('set'+str(ii)))
            strans = doc.createElement('strans')
            strans.setAttribute('reflection','0')
            strans.setAttribute('absolutemagnification','0')
            strans.setAttribute('absoluteangle','0')
            xy = doc.createElement('xy')
            xy.appendChild(doc.createTextNode('0,0'))

            sref.appendChild(sname)
            sref.appendChild(strans)
            sref.appendChild(xy)
            bgnstr.appendChild(sref)
        biginlib.appendChild(bgnstr)
        f=open(filename+'.xml', 'w')
        f.write(doc.toprettyxml(encoding='Shift_JIS'))

class grid(field):
    def __init__(self,fieldsize,nx,ny):
        field.__init__(self,fieldsize)
        self.nx            = nx
        self.ny            = ny
        self.markerwidth   = 1000
        self.markerlengths = [10000,5000]
        self.tickwidth     =  500
        self.ticklength    = 1000
        self.spacing = point(math.floor(self.fieldsize/self.nx),math.floor(self.fieldsize/self.ny))

        shapes=[]
        for i in range(self.nx+1):
            for j in range(self.ny+1):
                # draw top-left marker
                if i==0 and j==0:
                    shapes.append(cross(self.markerlengths,self.markerwidth,[True, False, False, True]).translate(self.markerwidth/2.0,-self.markerwidth/2.0))
                else:
                    drawarms = [True, True, True, True]
                    widths   = [self.tickwidth, self.tickwidth, self.tickwidth, self.tickwidth]
                    offset   = [0, 0]
                    if i==0:
                        drawarms[2] = False
                        widths[1]   = self.tickwidth/2.0
                        widths[3]   = self.tickwidth/2.0
                        offset[0]   = self.tickwidth/4.0
                    elif i==self.nx:
                        drawarms[0] = False
                        widths[1]   = self.tickwidth/2.0
                        widths[3]   = self.tickwidth/2.0
                        offset[0]   = -self.tickwidth/4.0
                    if j==0:
                        drawarms[1] = False
                        widths[0]   = self.tickwidth/2.0
                        widths[2]   = self.tickwidth/2.0
                        offset[1]   = -self.tickwidth/4.0
                    elif j==self.ny:
                        drawarms[3] = False
                        widths[0]   = self.tickwidth/2.0
                        widths[2]   = self.tickwidth/2.0
                        offset[1]   = self.tickwidth/4.0

                    nextshape=cross(self.ticklength,widths,drawarms)
                    nextshape.translate(self.spacing.x*i+offset[0],-self.spacing.y*j+offset[1])
                    shapes.append(nextshape)

        self.shapelists.append(shapes)
